import os
from datetime import datetime
from os.path import join

import torch
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

from ALE.data.dataset_gzsl import Dataset
from ALE.flags import parser
from ALE.models.confg_model import config_model
from common_utils.Evaluator import zsl_acc_gzsl2, zsl_acc_gzsl, zsl_acc2
from common_utils.utils import load_args

device = 'cuda' if torch.cuda.is_available() else 'cpu'

parser.add_argument('--phase', default='gzsql_test', help='训练阶段')
def main():
    print(os.getcwd())
    args = parser.parse_args()
    load_args(args.config, args)
    TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())
    logpath = os.path.join(args.cv_dir, args.name, 'tensorboard/'+args.phase+"/"+TIMESTAMP)
    os.makedirs(logpath, exist_ok=True)
    testSet = Dataset(data_dir=join(args.data_root,args.data_dir),
                      dataset=args.dataset, phase=args.phase)
    writer = SummaryWriter(log_dir= logpath, flush_secs=30)
    model, optimizer = config_model(args = args,feat_dim=testSet.feat_dim,attr_dim=testSet.attr_dim)
    model.eval()
    modelPath = os.path.join(args.cv_dir, args.name)
    checkpoint = torch.load(join(modelPath, "checkpoint"),map_location= device)
    model.load_state_dict(checkpoint['model_state_dict'])
    attrs = torch.index_select(testSet.sig, dim=0, index=testSet.labels_test_seen);
    _, test_seen_preds = model(imgs=testSet.test_seen_data.to(device), attrs=attrs.to(device),labels = testSet.labels_test_seen.to(device),sig = testSet.sig.to(device))
    acc_seen_classes = zsl_acc_gzsl2(attr_preds=test_seen_preds.data.cpu(),
                                 y_true=testSet.labels_test_seen,
                                 classes=testSet.test_classes_seen,
                                 sig=testSet.sig)
    attrs = torch.index_select(testSet.sig, dim=0,
                          index=testSet.labels_test_unseen)
    _, test_unseen_preds = model(imgs=testSet.test_unseen_data.to(device), attrs=attrs.to(device),labels = testSet.labels_test_unseen.to(device),sig = testSet.sig.to(device))
    acc_unseen_classes = zsl_acc_gzsl2(attr_preds=test_unseen_preds.data.cpu(),
                                   y_true=testSet.labels_test_unseen,
                                   classes=testSet.test_classes_unseen,
                                   sig=testSet.sig)
    czsl_unseen_acc = zsl_acc2(attr_preds=test_unseen_preds.data.cpu(), y_true=testSet.labels_test_unseen.data.cpu(),
                         sig=testSet.test_attrs_unseen.data.cpu())
    attrs = torch.index_select(testSet.sig, dim=0,
                              index=testSet.labels_test)
    _, test_preds = model(imgs=testSet.data.to(device), attrs=attrs.to(device),labels = testSet.labels_test.to(device),sig = testSet.sig.to(device))
    acc_classes = zsl_acc_gzsl2(attr_preds=test_preds.data.cpu(),
                                       y_true=testSet.labels_test,
                                       classes=testSet.test_classes,
                                       sig=testSet.sig)
    # 调和平均值
    HM = 2 * acc_seen_classes * acc_unseen_classes / (acc_seen_classes + acc_unseen_classes)
    print(f'U:{acc_unseen_classes}; S:{acc_seen_classes};all:{acc_classes};HM{HM},czsl_unseen_acc{czsl_unseen_acc}')
if __name__ == '__main__':
    main()
